import faiss
from typing import List, Tuple, Optional, Union
import numpy as np
import os
import pickle
from retrivolve.rag.rag_embeddings import LocalEmbeddings

class FaissVecotrDB:
    def __init__(self, dim: int, index_type: str, embeddings_model: LocalEmbeddings, device):
        self.dim = dim
        self.index_type = index_type
        self.device = device
        self.embeddings_model = embeddings_model
        self.id_map = {} # internal Faiss IDs -> user-defined marks
        self.reverse_id_map = {} # user-defined marks -> internal Faiss IDs
        self.mark_to_key = {} # user-defined marks -> original key string
        self.next_id = 0 # Next available internal Faiss ID

        self.index = self._build_index(dim, index_type)
        if device == "cuda":
            res = faiss.StandardGpuResources()
            self.index = faiss.index_cpu_to_all_gpus(res, 0, self.index)
    
    def _build_index(self, dim, index_type):
        if index_type == "Flat":
            return faiss.IndexFlatL2(dim)
        elif index_type.startswith("IVF"):
            nlist = int(index_type[3:])
            quantizer = faiss.IndexFlatL2(dim)
            return faiss.IndexIVFFlat(quantizer, dim, nlist, faiss.METRIC_L2)
        elif index_type.startswith("HNSW"):
            M = int(index_type[4:])
            return faiss.IndexHNSWFlat(dim, M)
        else:
            raise ValueError(f"Unsupported index type: {index_type}")
    
    def add_vector(self, keys: List[str], marks: Optional[List[str]]):
        vectors = self.embeddings_model.get_embedding(queries=keys)
        self._add_vector(vectors.numpy().tolist(), marks, keys)

    def _add_vector(self, vectors: Union[List[List[float]], np.ndarray], marks: Optional[List[str]] = None, keys: Optional[List[str]] = None):
        if len(vectors) == 0:
            return
        vectors = np.array(vectors).astype('float32')
        if vectors.ndim == 1:
            vectors = vectors.reshape(1, -1)
        if vectors.shape[1] != self.dim:
            raise ValueError(f"Vector dimension {vectors.shape[1]} does not match index dimension {self.dim}")
        n = len(vectors)
        if marks is None:
            marks = [f'id_{self.next_id+i}' for i in range(n)]
        if keys is not None and len(keys) != n:
            raise ValueError(f"Length of keys ({len(keys)}) must match number of vectors ({n})")
        for i, _id in enumerate(marks):
            if _id in self.reverse_id_map:
                raise ValueError(f"Mark {_id} already exists.")
            self.id_map[self.next_id] = _id
            self.reverse_id_map[_id] = self.next_id
            if keys is not None:
                self.mark_to_key[_id] = keys[i]
            self.next_id += 1
        int_ids = np.array([self.reverse_id_map[mark] for mark in marks]).astype('int64')

        if self.index_type.startswith("IVF") and not self.index.is_trained:
            self.index.train(vectors)
        if self.index_type == "Flat":
            self.index.add(vectors)
        else:
            self.index.add_with_ids(vectors, int_ids)

    
    def search(self, queries:List[str], k=5):
        query_vector = self.embeddings_model.get_embedding(queries=queries)
        results = self._search(query_vector.numpy().tolist(), k)
        return results

    def _search(self, query_vector: Union[List[float],np.ndarray], k=5):
        if isinstance(query_vector, list):
            query_vector = np.array(query_vector).astype('float32')
        if query_vector.ndim == 1:
            query_vector = np.array(query_vector).astype('float32').reshape(1, -1)
        distances, indices = self.index.search(query_vector, k)
        results = []
        for distance,indice in zip(distances, indices):
            result = []
            for dist, idx in zip(distance,indice):
                if idx == -1:
                    result.append((None, float('inf')))
                else:
                    mark = self.id_map[idx]
                    key_str = self.mark_to_key.get(mark, None)
                    result.append((key_str, dist))
            results.append(result)
        return results
    
    def save(self, file_path: str):
        os.makedirs(file_path, exist_ok=True)
        faiss.write_index(self.index, os.path.join(file_path, "index.faiss"))
        with open(os.path.join(file_path, "id_map.pkl"), 'wb') as f:
            pickle.dump((self.id_map, self.reverse_id_map, self.next_id, self.dim, self.index_type, self.mark_to_key), f)
    
    def load(self, file_path: str):
        self.index = faiss.read_index(os.path.join(file_path, "index.faiss"))
        with open(os.path.join(file_path, "id_map.pkl"), 'rb') as f:
            loaded = pickle.load(f)
            if len(loaded) == 6:
                self.id_map, self.reverse_id_map, self.next_id, self.dim, self.index_type, self.mark_to_key = loaded
            else:
                # backward compatibility for older files without mark_to_key
                self.id_map, self.reverse_id_map, self.next_id, self.dim, self.index_type = loaded
                self.mark_to_key = {mark: None for mark in self.reverse_id_map.keys()}
        if self.device == "cuda":
            res = faiss.StandardGpuResources()
            self.index = faiss.index_cpu_to_all_gpus(res, 0, self.index)